// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#include <nebula/types/union_type.h>
#include <turbo/log/logging.h>

namespace nebula {

    constexpr Type::type SparseUnionType::type_id;

    constexpr Type::type DenseUnionType::type_id;




    // ----------------------------------------------------------------------
    // Union type

    constexpr int8_t UnionType::kMaxTypeCode;
    constexpr int UnionType::kInvalidChildId;

    UnionMode::type UnionType::mode() const {
        return id_ == Type::SPARSE_UNION ? UnionMode::SPARSE : UnionMode::DENSE;
    }

    UnionType::UnionType(FieldVector fields, std::vector<int8_t> type_codes, Type::type id)
            : NestedType(id),
              type_codes_(std::move(type_codes)),
              child_ids_(kMaxTypeCode + 1, kInvalidChildId) {
        children_ = std::move(fields);
        KCHECK_OK(ValidateParameters(children_, type_codes_, mode()));
        for (int child_id = 0; child_id < static_cast<int>(type_codes_.size()); ++child_id) {
            const auto type_code = type_codes_[child_id];
            child_ids_[type_code] = child_id;
        }
    }

    turbo::Status UnionType::ValidateParameters(const FieldVector &fields,
                                                const std::vector<int8_t> &type_codes,
                                                UnionMode::type mode) {
        if (fields.size() != type_codes.size()) {
            return turbo::invalid_argument_error("Union should get the same number of fields as type codes");
        }
        for (const auto type_code: type_codes) {
            if (type_code < 0 || type_code > kMaxTypeCode) {
                return turbo::invalid_argument_error("Union type code out of bounds");
            }
        }
        return turbo::OkStatus();
    }

    DataTypeLayout UnionType::layout() const {
        if (mode() == UnionMode::SPARSE) {
            return DataTypeLayout(
                    {DataTypeLayout::AlwaysNull(), DataTypeLayout::FixedWidth(sizeof(uint8_t))});
        } else {
            return DataTypeLayout({DataTypeLayout::AlwaysNull(),
                                   DataTypeLayout::FixedWidth(sizeof(uint8_t)),
                                   DataTypeLayout::FixedWidth(sizeof(int32_t))});
        }
    }

    uint8_t UnionType::max_type_code() const {
        return type_codes_.size() == 0
               ? 0
               : *std::max_element(type_codes_.begin(), type_codes_.end());
    }

    std::string UnionType::to_string(bool show_metadata) const {
        std::stringstream s;

        s << name() << "<";

        for (size_t i = 0; i < children_.size(); ++i) {
            if (i) {
                s << ", ";
            }
            s << children_[i]->to_string(show_metadata) << "=" << static_cast<int>(type_codes_[i]);
        }
        s << ">";
        return s.str();
    }

    SparseUnionType::SparseUnionType(FieldVector fields, std::vector<int8_t> type_codes)
            : UnionType(std::move(fields), std::move(type_codes), Type::SPARSE_UNION) {}

    turbo::Result<std::shared_ptr<DataType>> SparseUnionType::create(FieldVector fields,
                                                                     std::vector<int8_t> type_codes) {
        TURBO_RETURN_NOT_OK(ValidateParameters(fields, type_codes, UnionMode::SPARSE));
        return std::make_shared<SparseUnionType>(std::move(fields), std::move(type_codes));
    }

    DenseUnionType::DenseUnionType(FieldVector fields, std::vector<int8_t> type_codes)
            : UnionType(std::move(fields), std::move(type_codes), Type::DENSE_UNION) {}

    turbo::Result<std::shared_ptr<DataType>> DenseUnionType::create(FieldVector fields,
                                                                    std::vector<int8_t> type_codes) {
        TURBO_RETURN_NOT_OK(ValidateParameters(fields, type_codes, UnionMode::DENSE));
        return std::make_shared<DenseUnionType>(std::move(fields), std::move(type_codes));
    }

    std::string UnionType::compute_fingerprint() const {
        std::stringstream ss;
        ss << type_id_fingerprint(*this);
        switch (mode()) {
            case UnionMode::SPARSE:
                ss << "[s";
                break;
            case UnionMode::DENSE:
                ss << "[d";
                break;
            default:
                KCHECK(false) << "Unexpected UnionMode";
        }
        for (const auto code: type_codes_) {
            // Represent code as integer, not raw character
            ss << ':' << static_cast<int32_t>(code);
        }
        ss << "]{";
        for (const auto &child: children_) {
            const auto &child_fingerprint = child->fingerprint();
            if (child_fingerprint.empty()) {
                return "";
            }
            ss << child_fingerprint << ";";
        }
        ss << "}";
        return ss.str();
    }

}  // namespace nebula
